We should expose these functions within the tex.grouped_mlp_experimental submodule:

It would make more sense to organize:

csrc/
├── extensions/
│   ├── grouped_mlp_experimental/
│   │   ├── megacpp.cpp
│   │   └── grouped_mlp_experimental.cpp
│   ├── pybind.cpp
│   └── ...

If we implement more mega-C++ impls in the future, I don't see a reason why they would be more similar to each other than to the block they are fusing.

Is it worth supporting GLU interleaving in the mega-C++ path? The only benefit is to support the fused GEMM+GLU kernel, and otherwise the unnecessary memory-bound kernel means perf is a lost cause. If we can simplify our optimized code paths, then it's worth it.

The only benefit is to support the fused GEMM+GLU kernel

I do hope in the future we can launch CuteDSL fused kernels in C++ with some TVM-FFI tricks, otherwise we are forced to choose either better kernel fusions or less CPU overhead. Currently the CuteDSL fusion path is very CPU bounded for small models and we rely on CUDA graph and paged stashing for it to work well

Why not make a separate fused op for mega-C++ CuTe DSL? It'll make the implementations less entangled, so there are fewer edge cases or complications that an agent might misunderstand.

yes, if there is a cutedsl version with mega C++, it's gonna take its own code path with zero code reuse since it's gonna be agent-assisted coding anyway

The GEMM+act fusions provide better GPU perf, so I think they should take higher priority than mega-C++. Basically, I see mega-C++ as "we can't do any better on GPU than the unfused impl, but at least we can make the CPU overhead very small".

Current order is follows:

1. check env var
2. env var = 1, then check supported recipe for mega-C++, so bf16 is supported, not mxfp8 / nvfp4
3. then for mxfp8, nvfp4, mega-C++ does fallback and check for the next fusion.

The reasoning is that, I do not want the compromise of either better fusion or less host bound, so for future mxfp8 support, we can do the following two things:

1. directly do cuteDSL integration directly with tvm-ffi and do cublas as a backup plan
2. maybe add a new value to NVTE_MEGACPP_GROUPED_LINEAR=forced, so for users who cannot enable cuda graph for some reason, they can enforce C++ when they know that their training is more host bound

The if recipe.mxfp8() or recipe.nvfp4(): return False branch is dead code — any non-None recipe that is not mxfp8/nvfp4 falls through to the identical final return False. The function is equivalent to return recipe is None; the current form misleads readers about which cases are being explicitly handled and can confuse future engineers adding support for a new recipe type.

The @functools.lru_cache(maxsize=None) with a CUDA stream handle as part of the key creates an unbounded cache. Each unique (num_groups, device_index, stream_handle) triplet permanently holds three CUDA tensors. In pipeline-parallel configurations or test suites that construct many torch.cuda.Stream() objects over time, destroyed streams' allocations are never freed because lru_cache holds the only live reference. Bounding the cache (e.g., maxsize=64) caps the worst-case retained GPU memory.

input_requires_grad is set to the generic requires_grad flag, which is True whenever any of the three op contexts requires a gradient — including the weight-only case. When the input tensor is frozen, the C++ backward still computes the full FC1 dgrad GEMM and discards the result. Using fc1_ctx.input_requires_grad as an additional gate matches the existing fuser convention and avoids the wasted computation.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

delay_wgrad_compute=True incompatibility is detected inside fuser_backward, which runs during .backward() — after the forward pass has already executed. A user who constructs a megacpp-fused model with delay_wgrad_compute=True will get through an entire forward step before hitting the ValueError. Moving this check to ForwardGroupedMLP_MegaCpp.__init__ or fuse_forward_megacpp_ops would surface the error at model-construction or fusion time instead.